// (c) Copyright slimCODE Software Inc. - www.slimcode.com
// This source is subject to the Microsoft Public License (Ms-PL).
// Please see http://go.microsoft.com/fwlink/?LinkID=131993 for details.
// All other rights reserved.

using System;
using System.Collections.Generic;
using System.IO;
using System.Security.Cryptography;
using System.Text;

namespace SlimCode.Utils
{
  internal class AesStream : Stream
  {
    public AesStream( Stream innerStream, string password, int strength, bool writing )
    {
      if( innerStream == null )
        throw new ArgumentException( "innerStream" );

      if( writing && !innerStream.CanWrite )
        throw new ArgumentException( "The inner stream cannot be written to.", "innerStream" );

      if( !writing && !innerStream.CanRead )
        throw new ArgumentException( "The inner stream cannot be read from.", "innerStream" );

      if( password == null )
        throw new ArgumentNullException( "password" );

      if( password.Length == 0 )
        throw new ArgumentException( "The password cannot be empty.", "password" );

      if( strength < 128 )
        throw new ArgumentException( "The strength cannot be less than 128." );

      if( ( strength % 64 ) != 0 )
        throw new ArgumentException( "The strength must be a factor of 64, not less than 128.", "strength" );

      int saltLength = strength / 16;
      byte[] salt = new byte[ saltLength ];

      if( writing )
      {
        // Generate salt.
        RandomNumberGenerator random = RandomNumberGenerator.Create();
        random.GetBytes( salt );

        // Take the opportunity to write salt.
        innerStream.Write( salt, 0, salt.Length );
      }
      else
      {
        // The inner stream should contain the salt.
        int total = 0;

        while( total < salt.Length )
        {
          int read = innerStream.Read( salt, total, salt.Length - total );

          if( read == 0 )
            throw new IOException( "Unexpected end of stream while reading the initialization vector." );

          total += read;
        }
      }

      Rfc2898DeriveBytes deriveBytes = new Rfc2898DeriveBytes( password, salt );

      byte[] encryptionKey = deriveBytes.GetBytes( saltLength * 2 );
      byte[] authenticationKey = deriveBytes.GetBytes( saltLength * 2 );
      byte[] verifierKey = deriveBytes.GetBytes( 2 );

      if( writing )
      {
        // Write the verifier.
        innerStream.Write( verifierKey, 0, verifierKey.Length );
      }
      else
      {
        // Compare the verifiers
        byte[] verifierRead = new byte[ 2 ];

        int read = innerStream.Read( verifierRead, 0, 2 );

        if( read == 1 )
        {
          read = innerStream.Read( verifierRead, 1, 1 );
        }

        if( read == 0 )
          throw new IOException( "Unexpected end of stream while reading the initialization vector." );

        if( ( verifierKey[ 0 ] != verifierRead[ 0 ] )
          || ( verifierKey[ 1 ] != verifierRead[ 1 ] ) )
          throw new CryptographicException( "Invalid password." );
      }

      m_rijndael = new RijndaelCtrMode( encryptionKey, 16 );
      m_innerStream = new AuthenticationStream( innerStream, authenticationKey, AuthenticationType.SHA512, writing );
      m_writing = writing;
    }

    public override bool CanRead
    {
      get { return !m_writing; }
    }

    public override bool CanSeek
    {
      get { return false; }
    }

    public override bool CanWrite
    {
      get { return m_writing; }
    }

    public override long Length
    {
      get { return m_innerStream.Length; }
    }

    public override long Position
    {
      get { return m_innerStream.Position; }
      set { throw new NotSupportedException( "This stream does not support seeking." ); }
    }

    public override void Flush()
    {
      m_innerStream.Flush();
    }

    public override int Read( byte[] buffer, int offset, int count )
    {
      if( m_writing )
        throw new IOException( "Cannot read from this stream. It is open for writing." );

      if( buffer == null )
        throw new ArgumentNullException( "buffer" );

      if( ( offset < 0 ) || ( offset >= buffer.Length ) )
        throw new ArgumentOutOfRangeException( "offset" );

      if( count < 0 )
        throw new ArgumentOutOfRangeException( "count" );

      if( offset + count > buffer.Length )
        throw new ArgumentOutOfRangeException( "count" );

      if( count == 0 )
        return 0;

      int read = m_innerStream.Read( buffer, offset, count );

      if( read > 0 )
      {
        m_rijndael.Encrypt( buffer, offset, read );
      }

      return read;
    }

    public override long Seek( long offset, SeekOrigin origin )
    {
      throw new NotSupportedException( "This stream does not support seeking." );
    }

    public override void SetLength( long value )
    {
      throw new NotSupportedException( "This stream does not support seeking." );
    }

    public override void Write( byte[] buffer, int offset, int count )
    {
      if( !m_writing )
        throw new IOException( "Cannot write to this stream. It is open for reading." );

      if( buffer == null )
        throw new ArgumentNullException( "buffer" );

      if( ( offset < 0 ) || ( offset >= buffer.Length ) )
        throw new ArgumentOutOfRangeException( "offset" );

      if( count < 0 )
        throw new ArgumentOutOfRangeException( "count" );

      if( offset + count > buffer.Length )
        throw new ArgumentOutOfRangeException( "count" );

      if( count > 0 )
      {
        m_rijndael.Encrypt( buffer, offset, count );
      }

      m_innerStream.Write( buffer, offset, count );
    }

    protected override void Dispose( bool disposing )
    {
      try
      {
        if( disposing )
        {
          if( m_innerStream != null )
          {
            m_innerStream.Close();
          }

          if( m_rijndael != null )
          {
            m_rijndael.Dispose();
          }
        }

        m_innerStream = null;
        m_rijndael = null;
      }
      finally
      {
        base.Dispose( disposing );
      }
    }

    private Stream m_innerStream; // = null
    private RijndaelCtrMode m_rijndael; // = null
    private bool m_writing; // = false
  }
}
